/* *******************************************************************
 * Copyright (c) 2004 IBM Corporation All rights reserved. This program and the accompanying materials are made
 * available under the terms of the Eclipse Public License v1.0 which accompanies this distribution and is available at
 * http://www.eclipse.org/legal/epl-v10.html
 * 
 * Contributors: Matthew Webster, Adrian Colyer, Martin Lippert initial implementation Andy Clement Abraham Nevado
 * *****************************************************************
 */
package org.mimicry.bridge.weaving;

import java.io.IOException;
import java.net.URL;
import java.security.CodeSource;
import java.util.HashMap;
import java.util.Map;

import org.aspectj.bridge.AbortException;
import org.aspectj.weaver.bcel.ExtensibleURLClassLoader;
import org.aspectj.weaver.loadtime.ClassLoaderWeavingAdaptor;
import org.aspectj.weaver.loadtime.DefaultWeavingContext;
import org.aspectj.weaver.loadtime.WeavingURLClassLoader;
import org.aspectj.weaver.tools.Trace;
import org.aspectj.weaver.tools.TraceFactory;
import org.aspectj.weaver.tools.WeavingAdaptor;
import org.aspectj.weaver.tools.WeavingClassLoader;

public class PostProcessingEnabledWeavingURLClassLoader extends ExtensibleURLClassLoader implements WeavingClassLoader
{
    private final URL[] aspectURLs;
    private WeavingAdaptor adaptor;
    private boolean initializingAdaptor;
    private final Map<String, byte[]> generatedClasses = new HashMap<String, byte[]>();

    private static Trace trace = TraceFactory.getTraceFactory().getTrace(WeavingURLClassLoader.class);

    public PostProcessingEnabledWeavingURLClassLoader(URL[] classURLs, URL[] aspectURLs, ClassLoader parent)
    {
        super(classURLs, parent);
        this.aspectURLs = aspectURLs;

        /*
         * If either we nor our parent is using an ASPECT_PATH use a new-style adaptor
         */
        if (this.aspectURLs.length > 0 || getParent() instanceof WeavingClassLoader)
        {
            try
            {
                adaptor = new WeavingAdaptor(this);
            }
            catch (ExceptionInInitializerError ex)
            {
                ex.printStackTrace(System.out);
                throw ex;
            }
        }
    }

    @Override
    protected void addURL(URL url)
    {
        if (adaptor == null)
        {
            createAdaptor();
        }
        adaptor.addURL(url);
        super.addURL(url);
    }

    protected byte[] postProcess(String name, byte[] byteCode)
    {
        return byteCode;
    }

    /**
     * Override to weave class using WeavingAdaptor
     */
    @Override
    protected Class<?> defineClass(String name, byte[] b, CodeSource cs) throws IOException
    {
        if (trace.isTraceEnabled())
        {
            trace.enter("defineClass", this, new Object[] { name, b, cs });
        }
        byte orig[] = b;
        /* Avoid recursion during adaptor initialization */
        if (!initializingAdaptor)
        {

            /* Need to defer creation because of possible recursion during constructor execution */
            if (adaptor == null && !initializingAdaptor)
            {
                createAdaptor();
            }

            try
            {
                b = adaptor.weaveClass(name, b, false);
            }
            catch (AbortException ex)
            {
                trace.error("defineClass", ex);
                throw ex;
            }
            catch (Throwable th)
            {
                trace.error("defineClass", th);
            }
        }

        b = postProcess(name, b);

        Class<?> clazz;
        // On error, define the original form of the class and log the issue
        try
        {
            clazz = super.defineClass(name, b, cs);
        }
        catch (Throwable th)
        {
            trace.error("Weaving class problem. Original class has been returned. The error was caused because of: "
                    + th, th);
            clazz = super.defineClass(name, orig, cs);
        }
        if (trace.isTraceEnabled())
        {
            trace.exit("defineClass", clazz);
        }
        return clazz;
    }

    private void createAdaptor()
    {
        DefaultWeavingContext weavingContext = new DefaultWeavingContext(this)
        {

            /* Ensures consistent LTW messages for testing */
            @Override
            public String getClassLoaderName()
            {
                ClassLoader loader = getClassLoader();
                return loader.getClass().getName();
            }

        };

        ClassLoaderWeavingAdaptor clwAdaptor = new ClassLoaderWeavingAdaptor();
        initializingAdaptor = true;
        clwAdaptor.initialize(this, weavingContext);
        initializingAdaptor = false;
        adaptor = clwAdaptor;
    }

    /**
     * Override to find classes generated by WeavingAdaptor
     */
    @Override
    protected byte[] getBytes(String name) throws IOException
    {
        byte[] bytes = super.getBytes(name);

        if (bytes == null)
        {
            return generatedClasses.remove(name);
        }

        return bytes;
    }

    /**
     * Implement method from WeavingClassLoader
     */
    @Override
    public URL[] getAspectURLs()
    {
        return aspectURLs;
    }

    @Override
    public void acceptClass(String name, byte[] classBytes, byte[] weavedBytes)
    {
        generatedClasses.put(name, weavedBytes);
    }
}
